/**
 * @file AMG.h
 * @author shanyoumu
 * @brief 包含一个完整的代数多重网格的实现，
 *        其中投影算子以及插值算子都是通过矩阵形式给出
 * 
 * @version 0.1
 * @date 2022-07-09
 * 
 * @copyright Copyright (c) 2022
 * 
 */

#ifndef AMG_H_
#define AMG_H_

#include <cmath>
#include <iostream>
#include <vector>
#include <map>
#include <functional>
#include <chrono>
#include <eigen3/Eigen/Sparse>
#include <eigen3/Eigen/IterativeLinearSolvers>
#include <eigen3/Eigen/Dense>
using std::vector;
using std::function;
using std::pair;
using SpMat = Eigen::SparseMatrix<double>;
using Tri = Eigen::Triplet<double>;
using u32 = unsigned int;
using Eigen::VectorXd;
using std::cout;
using std::endl;

/**
 * @brief 因为代数多重网格涉及到了单点的修改并求全局最大值，
 * 所以采用线段树的数据结构，以尽快找出最大值
 * 
 */
class SegmentTree 
{
private:
    vector<int> nums;        //第一个是点的编号，第二个是集合的势
    vector<pair<u32, int>> tree;
    u32 length;
    void pushUp(u32 i)
    {
        tree[i] = tree[2 * i + 1];
        if (tree[2 * i].second > tree[2 * i + 1].second)
            tree[i] = tree[2 * i];    
    }
    void build(u32 s, u32 t, u32 i)
    {
        if (s == t){
            tree[i].second = nums[s];
            tree[i].first = s;
            return;
        }
        u32 c = s + (t - s) / 2;
        build(s, c, 2 * i);
        build(c + 1, t, 2 * i + 1);
        pushUp(i);
    }
    void update(u32 idx, int x, u32 s, u32 t, u32 i)
    {
        if (s == t)
        {
            tree[i].second = x;
            return;
        }
        u32 c = s + (t - s) / 2;
        if (idx <= c)
            update(idx, x, s, c, 2 * i);
        else
            update(idx, x, c + 1, t, 2 * i + 1);
        pushUp(i);
    }
    void add(u32 idx, int x, u32 s, u32 t, u32 i)
    {
        if (s == t)
        {
            tree[i].second += x;
            return;
        }
        u32 c = s + (t - s) / 2;
        if (idx <= c)
            add(idx, x, s, c, 2 * i);
        else
            add(idx, x, c + 1, t, 2 * i + 1);
        pushUp(i);
    }
public:
    SegmentTree(vector<int> &_nums)
    {
        nums = _nums;
        length = _nums.size();
        tree.resize(4 * length);
        build(0, length - 1, 1);
    }
    void update(u32 i, int x)       //将第i个位置的值改成x
    {
        update(i, x, 0, length - 1, 1);
    }
    void add(u32 i, int x)
    {
        add(i, x, 0, length - 1, 1);
    }
    pair<u32, int> getmax() {return tree[1];};
};

/**
 * @brief 所采用的四种迭代法
 * 
 */
enum InteratorMethon {WJ, GS, SOR, SSOR};

/**vector<pair<u32, double>> Boundaryval
 * @brief 返回两个向量之间的L2范数
 * 
 * @param v1 
 * @param v2 
 * @return double 
 */
double L2_Error(VectorXd v1, VectorXd v2); 

/**
 * @brief 对方程 Au=Rhs 以初始值u进行一次迭代，结果仍保留在u中
 * 
 * @param _A 
 * @param _u 
 * @param _Rhs 
 * @param _type 迭代法类型
 * @param weight 加权的大小
 */
void Interator(SpMat &_A, VectorXd &_u, VectorXd &_Rhs, InteratorMethon _type, double weight = 2.0/3);

class BuildAMG;

class BuildAMG
{
private:
    SpMat A;
    double theta;
    SpMat ItplatMat;        //保存最终得到的一个插值算子
    vector<double> beta;    //保存每一行非对角元的最大值乘以theta的值
    vector<u32> Cpoints;
public:
    BuildAMG() = delete;
    BuildAMG(SpMat &Spt, double _theta = 0.25) {A = Spt; theta = _theta;};
    /// @brief 注意，这里把核心点选取与插值算子的生成分开，因为插值算子的生成可能会多种方式，核心点选取是固定的
    /// 当然也可以自己输入然后进行下一步的插值矩阵的生成
    void PreBuild();        //设置beta和Cpoints,从矩阵中得到
    void setCpoints(vector<u32> &cp) {Cpoints = cp;};   //从外界输入核心点
    void buildall();        //得到ItplatMat
    void Build() {PreBuild(); buildall();};
    SpMat GetItplatMat() {return ItplatMat;};
    vector<u32>& getcorepnt() {return Cpoints;};
};

class AMG_Multigrid;

class AMG_Multigrid
{
private:
    vector<SpMat> Vec_Mat;      //保存所有的矩阵，Vec_Mat[0]为初始网格，Vec_Mat[n-1]为最粗网格
    vector<VectorXd> F_Rhs;     //保存方程右端项
    vector<VectorXd> u_Lhs;     //保存每一层网格方程的近似解
    vector<SpMat> Interpolation;     //保存每一层网格的插值算子
    vector<SpMat> Restrict;            //保存每一层的投影算子
    Eigen::PartialPivLU<Eigen::MatrixXd> LastmatLU;     //保存最粗网格矩阵的LU分解
    // Eigen::MatrixXd Lastmatinv;		//保存最后一层矩阵的逆
    double Tolerance = 1e-13;           //求解使用的迭代法的终止精度
    size_t iter_num;            //迭代终止时整个过程所迭代的次数
    InteratorMethon InteratorType;      //迭代的方法
    double InteratorWeight;             //迭代的权重
    bool Has_Solved = false;            //判断是否已经解决
    int N_Layer;                        //记录网格层数
    double Theta = 0.25;                 //粗化网格的参数
public:
    AMG_Multigrid() = delete;
    /// @brief 构造代数多重网格
    /// @param _Spt 待求解方程的左端项
    /// @param _Rhs 待求解方程的右端项
    /// @param _Tolerance 设置计算阈值（阈值实际为该值乘以初始误差）
    /// @param _InteratorType 迭代方法（WJ,GS,SOR,SSOR）
    /// @param _InteratorWeight 迭代法的权重
    /// @param theta 代数多重网格的强连接点选取的阈值
    AMG_Multigrid(SpMat &_Spt, VectorXd &_Rhs, double _Tolerance = 1e-8, InteratorMethon _InteratorType = GS, double _InteratorWeight = 2.0 / 3, double theta = 0.25);             //初始化所有的投影矩阵，网格矩阵等
    void Set_InteratorType(InteratorMethon _InteratorType, double _Weight);     //设置迭代法
    void Set_Tolerance(double _Tolerance);      //设置求解精度
    void V_Cycle();                 //用Vcycle解方程
    void V_CycleOnce(int _NowLayer);             //做一次Vcycle
    void Solve() {cout << "Using AMG Solver :" << endl; V_Cycle(); };          //求解
    VectorXd GetResult() {return u_Lhs[0];};
    vector<SpMat> &Getmat() {return Vec_Mat;};
    vector<SpMat> &Getplt() {return Interpolation;};
};
#endif
